import UIKit
import Metal
import MetalKit

class MetalMp4Processor {
    
    // Metal相关属性
    var device: MTLDevice!
    private var commandQueue: MTLCommandQueue!
    private var pipelineState: MTLRenderPipelineState!
    private var vertexBuffer: MTLBuffer!
    var sampler: MTLSamplerState! // 纹理采样器
    
    
    init() {
        // 获取Metal设备
        guard let device = MTLCreateSystemDefaultDevice() else {
            fatalError("Metal is not supported on this device")
        }
        self.device = device
    
        
        // 创建命令队列
        commandQueue = device.makeCommandQueue()
        
        // 编译着色器
        let library = device.makeDefaultLibrary()
        let vertexFunction = library?.makeFunction(name: "vertexShader")
        let fragmentFunction = library?.makeFunction(name: "fragmentShader")
        
        // 配置渲染管线
        let pipelineDescriptor = MTLRenderPipelineDescriptor()
        pipelineDescriptor.vertexFunction = vertexFunction
        pipelineDescriptor.fragmentFunction = fragmentFunction
        pipelineDescriptor.colorAttachments[0].pixelFormat = .bgra8Unorm
        
        // 创建顶点描述符，定义顶点数据布局
        let vertexDescriptor = MTLVertexDescriptor()

        // 配置第一个属性（position：float4）
        vertexDescriptor.attributes[0].format = .float4    // 数据格式
        vertexDescriptor.attributes[0].offset = 0           // 在顶点结构体中的偏移量
        vertexDescriptor.attributes[0].bufferIndex = 0      // 关联的缓冲区索引

        // 配置第二个属性（texCoord：float2）
        vertexDescriptor.attributes[1].format = .float2
        vertexDescriptor.attributes[1].offset = MemoryLayout<Float>.stride * 4  // 紧跟在position之后
        vertexDescriptor.attributes[1].bufferIndex = 0

        // 配置缓冲区布局（单个顶点的总大小）
        vertexDescriptor.layouts[0].stride = MemoryLayout<Float>.stride * 6
        vertexDescriptor.layouts[0].stepFunction = .perVertex  // 每个顶点对应一个数据
        vertexDescriptor.layouts[0].stepRate = 1
        
        // 关键：设置顶点描述符（解决stage_in错误的核心）
        pipelineDescriptor.vertexDescriptor = vertexDescriptor
        
        // 关键：开启透明混合（Alpha通道生效的核心配置）
        pipelineDescriptor.colorAttachments[0].isBlendingEnabled = true
        pipelineDescriptor.colorAttachments[0].rgbBlendOperation = .add // RGB混合方式：相加
        pipelineDescriptor.colorAttachments[0].alphaBlendOperation = .add // Alpha混合方式：相加
        // 混合因子：符合常规透明逻辑（源像素Alpha控制透明度）
        pipelineDescriptor.colorAttachments[0].sourceRGBBlendFactor = .sourceAlpha
        pipelineDescriptor.colorAttachments[0].sourceAlphaBlendFactor = .sourceAlpha
        pipelineDescriptor.colorAttachments[0].destinationRGBBlendFactor = .oneMinusSourceAlpha
        pipelineDescriptor.colorAttachments[0].destinationAlphaBlendFactor = .oneMinusSourceAlpha
        
        // 创建管线状态
        pipelineState = try! device.makeRenderPipelineState(descriptor: pipelineDescriptor)
        
        // 3. 创建纹理采样器（基础配置，确保纹理正常采样）
        let samplerDesc = MTLSamplerDescriptor()
        samplerDesc.minFilter = .linear // 缩小平滑
        samplerDesc.magFilter = .linear // 放大平滑
        samplerDesc.sAddressMode = .clampToEdge // 超出范围夹紧到边缘
        samplerDesc.tAddressMode = .clampToEdge
        self.sampler = device.makeSamplerState(descriptor: samplerDesc)!
        
        // 设置顶点数据（全屏矩形）
        let vertices: [Float] = [
            // 位置(x,y,z,w)      // 纹理坐标(u,v)
            -1.0, -1.0, 0.0, 1.0,  0.0, 1.0,
             1.0, -1.0, 0.0, 1.0,  1.0, 1.0,
            -1.0,  1.0, 0.0, 1.0,  0.0, 0.0,
             1.0,  1.0, 0.0, 1.0,  1.0, 0.0
        ]
        vertexBuffer = device.makeBuffer(bytes: vertices, length: vertices.count * MemoryLayout<Float>.stride, options: [])
    }
    
    // 从合并图中处理并生成原图大小的透明图
    func process(mtlTexture: MTLTexture, drawable: CAMetalDrawable) {

        let renderPassDescriptor = MTLRenderPassDescriptor()
        renderPassDescriptor.colorAttachments[0].texture = drawable.texture
        renderPassDescriptor.colorAttachments[0].loadAction = .clear
        renderPassDescriptor.colorAttachments[0].clearColor = MTLClearColor(red: 0, green: 0, blue: 0, alpha: 0)
        renderPassDescriptor.colorAttachments[0].storeAction = .store
        
        guard let commandBuffer = commandQueue.makeCommandBuffer() else { return }
        guard let encoder = commandBuffer.makeRenderCommandEncoder(descriptor: renderPassDescriptor) else { return }
        
        encoder.setRenderPipelineState(pipelineState)
        encoder.setVertexBuffer(vertexBuffer, offset: 0, index: 0)
        // 绑定纹理和采样器（texture(0)、sampler(0)，对应Shader中的参数）
        encoder.setFragmentTexture(mtlTexture, index: 0)
        encoder.setFragmentSamplerState(sampler, index: 0)
        
        encoder.drawPrimitives(type: .triangleStrip, vertexStart: 0, vertexCount: 4)
        
        encoder.endEncoding()
        commandBuffer.present(drawable)
        commandBuffer.commit()
        
    }
    
}


    
